﻿/*  CLR版本: 4.0.30319.18063
 * 系统时间: 2014/10/31 14:20:43
 * 创建年份: 2014
 *     作者: 程炜.Snail
 */
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Security.Cryptography;
using System.Net.NetworkInformation;
using System.Net;
using System.Threading;
using System.Text.RegularExpressions;

namespace Quick.Communicator
{
    /// <summary>
    /// 基于SuperSocket 的WebSocket
    /// </summary>
    public class WebSocket
    {
        #region 属性
        /// <summary>
        /// WebSocket 使用的TCP
        /// </summary>
        private TCPServer m_TcpServer = null;
        /// <summary>
        /// WebSocket 服务器握手协议中的常量GUID
        /// </summary>
        private const string m_Guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
        /// <summary>
        /// 是否已完成握手协议
        /// </summary>
        Dictionary<EndPoint, bool> m_HandshakeDictionary = new Dictionary<EndPoint, bool>();
        /// <summary>
        /// 线程锁
        /// </summary>
        private ReaderWriterLockSlim m_HandshakeRWLS = new ReaderWriterLockSlim();
        /// <summary>
        /// 当前服务器是否运行
        /// </summary>
        public bool IsRunning
        {
            get { return m_TcpServer.IsRunning; }
        }
        /// <summary>
        /// 当前服务器链接的客户端数量
        /// </summary>
        public int SessionCount
        {
            get { return m_TcpServer.SessionCount; }
        }
        /// <summary>
        /// 客户端链接 Session
        /// </summary>
        public Dictionary<EndPoint, Session> ClientSession
        {
            get { return m_TcpServer.ClientSession; }
        }
        /// <summary>
        /// 系统使用的数据报文
        /// </summary>
        public IResolver Resolver
        {
            set { m_TcpServer.Resolver = value; }
            get { return m_TcpServer.Resolver; }
        }
        #endregion

        #region 事件
        /// <summary>
        /// 客户端建立连接事件
        /// </summary>
        public event NetEventHandler ClientConn;
        /// <summary>
        /// 客户端关闭事件
        /// </summary>
        public event NetEventHandler ClientClose;
        /// <summary>
        /// 服务器配置变更
        /// </summary>
        public event NetEventHandler ServerConfigChange;
        /// <summary>
        /// 信号接收完成 事件
        /// </summary>
        public event ReceiveDataEndEventHandler ReceiveDataEnd;
        /// <summary>
        /// 信号发送完成 事件
        /// </summary>
        public event SendDataEndEventHandler SendDataEnd;
        /// <summary>
        /// 连接客户端数量发送改变
        /// </summary>
        public event ClientConnectionChangeEventHandler ClientConnectionChange;
        /// <summary>
        /// 服务器接收的客户端字节总数
        /// </summary>
        public event ReceiveTotalBytesEventHandler ReceiveTotalBytes;
        /// <summary>
        /// 是否禁止的IP
        /// </summary>
        public IsBanIPEventHandler IsBanIP;
        #endregion

        #region 构造函数
        /// <summary>
        /// 构造函数(默认使用ASCII编码方式)
        /// </summary>
        /// <param name="port">服务器端监听的端口号</param>
        /// <param name="timeouts">链接超时时间(分钟)</param>
        ///  <param name="maxClient">服务器能容纳客户端的最大能力</param>
        public WebSocket(int port, int timeouts = 10, int maxClient = 5000)
        {
            DatagramResolver resolver = new DatagramResolver(Encoding.UTF8);
            m_TcpServer = new TCPServer(port, maxClient, timeouts, resolver);
            m_TcpServer.ReceiveDataEnd += m_TcpServer_ReceiveDataEnd;
            m_TcpServer.SendDataEnd += OnSendCommandEnd;
            m_TcpServer.ClientConn += OnClientConn;
            m_TcpServer.ServerConfigChange += OnServerConfigChange;
            m_TcpServer.ClientConnectionChange += OnClientConnectionChange;
            m_TcpServer.ReceiveTotalBytes += OnReceiveTotalBytes;
        }
        #endregion

        #region WebSocket 控制
        /// <summary>
        /// WebSocket 启动监听
        /// </summary>
        public void Start()
        {
            m_TcpServer.Start();
        }
        /// <summary>
        /// 停止监听端口
        /// </summary>
        public void Stop()
        {
            m_TcpServer.Stop();
        }
        /// <summary>
        /// 重新设置端口号,超时时间,报文解析器
        /// </summary>
        /// <param name="port">监听端口号</param>
        /// <param name="Timeouts">超时时间</param>
        public void SetSocketParameter(int port, int timeouts)
        {
            m_TcpServer.SetSocketParameter(port, timeouts, new DatagramResolver(Encoding.UTF8));
        }
        #endregion

        #region 发送数据
        /// <summary>
        /// 服务器端发送数据的方法
        /// 如有设置报文结束符,则会添加报文结束符
        /// </summary>
        /// <param name="endPoint">数据的IP地址</param>
        /// <param name="datagram">需发送的报文数据(不包含报文结束符)</param>
        /// <param name="isCloseSession">发送完成后进行断开链接</param>
        public virtual bool AsyncSend(EndPoint endPoint, string datagram, bool isCloseSession = false)
        {
            return m_TcpServer.AsyncSend(endPoint, PackData(datagram), false, isCloseSession);
        }
        /// <summary>
        /// 服务器端发送数据的方法
        /// </summary>
        /// <param name="endPoint"></param>
        /// <param name="bt"></param>
        /// <param name="isCloseSession">发送完成后进行断开链接</param>
        public virtual bool AsyncSend(EndPoint endPoint, byte[] bt, bool isCloseSession = false)
        {
            return m_TcpServer.AsyncSend(endPoint, PackData(bt), false, isCloseSession);
        }
        /// <summary>
        /// 服务器端发送数据的方法
        /// </summary>
        /// <param name="recvDataClient">接收数据的客户端会话（特定的客户）</param>
        /// <param name="datagram">需发送的报文数据(不包含报文结束符)</param>
        /// <param name="isCloseSession">发送完成后进行断开链接</param>
        public virtual bool AsyncSend(Session recvDataClient, string datagram, bool isCloseSession = false)
        {
            return m_TcpServer.AsyncSend(recvDataClient, PackData(datagram), false, isCloseSession);
        }
        /// <summary>
        /// 服务器端发送数据的方法
        /// </summary>
        /// <param name="recvDataClient">接收数据的客户端会话（特定的客户）</param>
        /// <param name="bt">数据报文(不包含报文结束符)</param>
        /// <param name="isCloseSession">发送完成后进行断开链接</param>
        public virtual bool AsyncSend(Session recvDataClient, byte[] bt, bool isCloseSession = false)
        {
            return m_TcpServer.AsyncSend(recvDataClient, PackData(bt), false, isCloseSession);
        }
        #endregion

        #region 事件触发
        /// <summary>
        /// 数据报文发送完成
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void OnSendCommandEnd(object sender, SendStateEventArgs e)
        {
            if (SendDataEnd != null)
            {
                Delegate[] delegateList = SendDataEnd.GetInvocationList();
                foreach (SendDataEndEventHandler handler in delegateList)
                {
                    handler.BeginInvoke(sender, e, null, null);
                }
            }
        }
        /// <summary>
        /// 新的客户端连接,发出通知
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void OnClientConn(object sender, NetEventArgs e)
        {
            try
            {
                m_HandshakeRWLS.TryEnterWriteLock(1000);
                if (m_HandshakeDictionary.ContainsKey(e.Client.ID))
                {   //重新链接,表示客户端断开过,需要重新进行握手协议
                    m_HandshakeDictionary[e.Client.ID] = false;
                }
                else
                {
                    m_HandshakeDictionary.Add(e.Client.ID, false);
                }
            }
            catch (Exception)
            { }
            finally
            {
                if (m_HandshakeRWLS.IsWriteLockHeld)
                {
                    m_HandshakeRWLS.ExitWriteLock();
                }
            }
            if (ClientConn != null)
            {
                Delegate[] delegateList = ClientConn.GetInvocationList();
                foreach (NetEventHandler handler in delegateList)
                {
                    handler.BeginInvoke(sender, e, null, null);
                }
            }
        }
        /// <summary>
        /// Client 断开连接
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void OnClientClose(object sender, NetEventArgs e)
        {
            try
            {
                m_HandshakeRWLS.TryEnterWriteLock(1000);
                if (m_HandshakeDictionary.ContainsKey(e.Client.ID))
                {   //重新链接,表示客户端断开过,需要重新进行握手协议
                    m_HandshakeDictionary.Remove(e.Client.ID);
                }
            }
            catch (Exception)
            { }
            finally
            {
                if (m_HandshakeRWLS.IsWriteLockHeld)
                {
                    m_HandshakeRWLS.ExitWriteLock();
                }
            }
            if (null != ClientClose)
            {
                Delegate[] delegateList = ClientClose.GetInvocationList();
                foreach (NetEventHandler handler in delegateList)
                {
                    handler.BeginInvoke(sender, e, null, null);
                }
            }
        }
        /// <summary>
        /// 信号接收完成 事件,通知相关订阅者
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void OnReceiveSignalEnd(object sender, NetEventArgs e)
        {
            if (null != ReceiveDataEnd)
            {
                Delegate[] delegateList = ReceiveDataEnd.GetInvocationList();
                foreach (ReceiveDataEndEventHandler handler in delegateList)
                {
                    handler.BeginInvoke(sender, e, null, null);
                }
            }
        }
        /// <summary>
        /// 配置变更
        /// </summary>
        private void OnServerConfigChange(object sender, NetEventArgs e)
        {
            if (null != ServerConfigChange)
            {
                Delegate[] delegateList = ServerConfigChange.GetInvocationList();
                foreach (NetEventHandler handler in delegateList)
                {
                    handler.BeginInvoke(sender, e, null, null);
                }
            }
        }
        /// <summary>
        /// Client 连接发送改变
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void OnClientConnectionChange(object sender, ConnectCountEventArgs e)
        {
            if (null != ClientConnectionChange)
            {
                Delegate[] delegateList = ClientConnectionChange.GetInvocationList();
                foreach (ClientConnectionChangeEventHandler handler in delegateList)
                {
                    handler.BeginInvoke(sender, e, null, null);
                }
            }
        }
        /// <summary>
        /// 服务器接收字节总数
        /// </summary>
        /// <param name="sender"></param>
        private void OnReceiveTotalBytes(object sender, ReceiveTotalBytesEventArgs e)
        {
            if (null != ReceiveTotalBytes)
            {
                Delegate[] delegateList = ReceiveTotalBytes.GetInvocationList();
                foreach (ReceiveTotalBytesEventHandler handler in delegateList)
                {
                    handler.BeginInvoke(sender, e, null, null);
                }
            }
        }
        #endregion

        #region WebSocket 协议相关
        /// <summary>
        /// Socket接收到数据报文,需进行解析
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void m_TcpServer_ReceiveDataEnd(object sender, NetEventArgs e)
        {
            bool isAccept = false;
            if (m_HandshakeDictionary.TryGetValue(e.Client.ID, out isAccept) && !isAccept)
            {   //未握手
                try
                {
                    isAccept = m_TcpServer.AsyncSend(e.Client.ID, PackHandShakeData(GetSecKeyAccetp(e.Client.DatagramByte)), false, false);
                    m_HandshakeRWLS.TryEnterWriteLock(1000);
                    if (m_HandshakeDictionary.ContainsKey(e.Client.ID))
                    {   //重新链接,表示客户端断开过,需要重新进行握手协议
                        m_HandshakeDictionary[e.Client.ID] = isAccept;
                    }
                    else
                    {
                        m_HandshakeDictionary.Add(e.Client.ID, isAccept);
                    }
                }
                catch (Exception)
                { }
                finally
                {
                    if (m_HandshakeRWLS.IsWriteLockHeld)
                    {
                        m_HandshakeRWLS.ExitWriteLock();
                    }
                }
            }
            else //解析数据内容
            {
                if (e.Client.DatagramByte?.Length > 0)
                {
                    if ((e.Client.DatagramByte[0] & 15) == 8)
                    { //链接关闭
                        m_TcpServer.CloseClientSocket(e.Client.ID);
                        return;
                    }
                }
                e.Client.Datagram = AnalyticData(e.Client.DatagramByte);
                if (!string.IsNullOrWhiteSpace(e.Client.Datagram))
                {
                    e.Client.DatagramByte = Encoding.UTF8.GetBytes(e.Client.Datagram);
                    OnReceiveSignalEnd(sender, e);
                }
            }
        }
        /// <summary>
        /// 打包握手信息
        /// </summary>
        /// <param name="secKeyAccept">Sec-WebSocket-Accept</param>
        /// <returns>数据包</returns>
        private static byte[] PackHandShakeData(string secKeyAccept)
        {
            StringBuilder responseBuilder = new StringBuilder();
            responseBuilder.Append("HTTP/1.1 101 Switching Protocols" + Environment.NewLine);
            responseBuilder.Append("Upgrade: websocket" + Environment.NewLine);
            responseBuilder.Append("Connection: Upgrade" + Environment.NewLine);
            responseBuilder.Append("Sec-WebSocket-Accept: " + secKeyAccept + Environment.NewLine + Environment.NewLine);
            //如果把上一行换成下面两行，才是thewebsocketprotocol-17协议，但居然握手不成功，目前仍没弄明白！
            //responseBuilder.Append("Sec-WebSocket-Accept: " + secKeyAccept + Environment.NewLine);
            //responseBuilder.Append("Sec-WebSocket-Protocol: chat" + Environment.NewLine);
            return Encoding.UTF8.GetBytes(responseBuilder.ToString());
        }
        /// <summary>
        /// 生成Sec-WebSocket-Accept
        /// </summary>
        /// <param name="handShakeText">客户端握手信息</param>
        /// <returns>Sec-WebSocket-Accept</returns>
        private static string GetSecKeyAccetp(byte[] handShakeBytes)
        {
            string handShakeText = Encoding.UTF8.GetString(handShakeBytes);
            string key = string.Empty;
            Regex r = new Regex(@"Sec\-WebSocket\-Key:(.*?)\r\n");
            Match m = r.Match(handShakeText);
            if (m.Groups.Count != 0)
            {
                key = Regex.Replace(m.Value, @"Sec\-WebSocket\-Key:(.*?)\r\n", "$1").Trim();
            }
            byte[] encryptionString = SHA1.Create().ComputeHash(Encoding.ASCII.GetBytes(key + m_Guid));
            return Convert.ToBase64String(encryptionString);
        }
        /// <summary>
        /// 打包服务器数据
        /// </summary>
        /// <param name="message">数据</param>
        /// <returns>数据包</returns>
        private static byte[] PackData(string message)
        {
            byte[] msgBytes = Encoding.UTF8.GetBytes(message);
            return PackData(msgBytes);
        }
        /// <summary>
        /// 打包需要从服务器发送的数据
        /// </summary>
        /// <param name="msgBytes"></param>
        /// <returns></returns>
        private static byte[] PackData(byte[] msgBytes)
        {
            byte[] contentBytes = null;
            if (msgBytes.Length < 126)
            {
                contentBytes = new byte[msgBytes.Length + 2];
                contentBytes[0] = 0x81;
                contentBytes[1] = (byte)msgBytes.Length;
                Array.Copy(msgBytes, 0, contentBytes, 2, msgBytes.Length);
            }
            else if (msgBytes.Length < 0xFFFF)
            {
                contentBytes = new byte[msgBytes.Length + 4];
                contentBytes[0] = 0x81;
                contentBytes[1] = 126;
                contentBytes[2] = (byte)(msgBytes.Length & 0xFF);
                contentBytes[3] = (byte)(msgBytes.Length >> 8 & 0xFF);
                Array.Copy(msgBytes, 0, contentBytes, 4, msgBytes.Length);
            }
            else
            {
                // 暂不处理超长内容  
            }
            return contentBytes;
        }
        /// <summary>
        /// 解析客户端数据包
        /// 参考协议 http://www.cnblogs.com/smark/archive/2012/11/26/2789812.html
        /// </summary>
        /// <param name="recBytes">服务器接收的数据包</param>
        /// <param name="recByteLength">有效数据长度</param>
        /// <returns></returns>
        private static string AnalyticData(byte[] recBytes)
        {
            if (null == recBytes || recBytes.Length < 2)
            {
                return string.Empty;
            }
            bool fin = (recBytes[0] & 0x80) == 0x80; // 1bit，1表示最后一帧  
            if (!fin)
            {
                return string.Empty;// 超过一帧暂不处理
            }
            bool mask_flag = (recBytes[1] & 0x80) == 0x80; // 是否包含掩码  
            int payload_len = recBytes[1] & 0x7F; // 数据长度  
            if (payload_len == 0)
            {
                return string.Empty;
            }
            byte[] masks = new byte[4];
            byte[] payload_data;

            if (payload_len == 126)
            {
                Array.Copy(recBytes, 4, masks, 0, 4);
                payload_len = (UInt16)(recBytes[2] << 8 | recBytes[3]);
                payload_data = new byte[payload_len];
                Array.Copy(recBytes, 8, payload_data, 0, payload_len);
            }
            else if (payload_len == 127)
            {
                Array.Copy(recBytes, 10, masks, 0, 4);
                byte[] uInt64Bytes = new byte[8];
                for (int i = 0; i < 8; i++)
                {
                    uInt64Bytes[i] = recBytes[9 - i];
                }
                UInt64 len = BitConverter.ToUInt64(uInt64Bytes, 0);

                payload_data = new byte[len];
                for (UInt64 i = 0; i < len; i++)
                {
                    payload_data[i] = recBytes[i + 14];
                }
            }
            else
            {
                Array.Copy(recBytes, 2, masks, 0, 4);
                payload_data = new byte[payload_len];
                Array.Copy(recBytes, 6, payload_data, 0, payload_len);
            }
            if (mask_flag) //有掩码,进行转换
            {
                for (var i = 0; i < payload_len; i++)
                {
                    payload_data[i] = (byte)(payload_data[i] ^ masks[i % 4]);
                }
            }
            return Encoding.UTF8.GetString(payload_data);
        }
        #endregion
    }
}
